from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
                            QLabel, QTextEdit, QPushButton, QLineEdit, QComboBox, 
                            QGroupBox, QTabWidget, QStatusBar, QMessageBox)
from PyQt5.QtCore import Qt
import sys
import json
import requests
from pathlib import Path

class CodeImproverGUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("DeepSeek 代码自主迭代系统")
        self.resize(1000, 800)
        
        # 主窗口中心部件
        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)
        
        # 主布局
        self.main_layout = QVBoxLayout(self.central_widget)
        
        # 创建API设置区域
        self.create_api_settings()
        
        # 创建代码编辑区域
        self.create_code_editor()
        
        # 创建控制按钮区域
        self.create_control_buttons()
        
        # 创建日志输出区域
        self.create_log_output()
        
        # 状态栏
        self.status_bar = QStatusBar()
        self.setStatusBar(self.status_bar)
        
        # 初始化API设置
        self.api_key = self.load_api_key()
        self.api_key_input.setText(self.api_key if self.api_key else "")
    
    def create_api_settings(self):
        api_group = QGroupBox("API 设置")
        api_layout = QHBoxLayout()
        
        api_group.setLayout(api_layout)
        self.main_layout.addWidget(api_group)
        
        # API Key输入
        api_layout.addWidget(QLabel("API Key:"))
        self.api_key_input = QLineEdit()
        self.api_key_input.setPlaceholderText("输入DeepSeek API Key")
        self.api_key_input.setEchoMode(QLineEdit.Password)
        api_layout.addWidget(self.api_key_input)
        
        # 保存按钮
        self.save_button = QPushButton("保存")
        self.save_button.clicked.connect(self.save_api_key)
        api_layout.addWidget(self.save_button)
        
        # 模型选择
        api_layout.addWidget(QLabel("模型:"))
        self.model_combo = QComboBox()
        self.model_combo.addItems(["deepseek-chat", "deepseek-reasoner"])
        api_layout.addWidget(self.model_combo)
    
    def create_code_editor(self):
        self.tab_widget = QTabWidget()
        self.main_layout.addWidget(self.tab_widget)
        
        # 原始代码标签页
        self.original_tab = QWidget()
        self.original_layout = QVBoxLayout()
        self.original_tab.setLayout(self.original_layout)
        
        self.original_code_edit = QTextEdit()
        self.original_layout.addWidget(QLabel("原始代码"))
        self.original_layout.addWidget(self.original_code_edit)
        
        # 需求标签页
        self.requirements_tab = QWidget()
        self.requirements_layout = QVBoxLayout()
        self.requirements_tab.setLayout(self.requirements_layout)
        
        self.requirements_edit = QTextEdit()
        self.requirements_layout.addWidget(QLabel("需求说明"))
        self.requirements_layout.addWidget(self.requirements_edit)
        
        # 测试用例标签页
        self.test_cases_tab = QWidget()
        self.test_cases_layout = QVBoxLayout()
        self.test_cases_tab.setLayout(self.test_cases_layout)
        
        self.test_cases_edit = QTextEdit()
        self.test_cases_layout.addWidget(QLabel("测试用例"))
        self.test_cases_layout.addWidget(self.test_cases_edit)
        
        # 改进后代码标签页
        self.improved_tab = QWidget()
        self.improved_layout = QVBoxLayout()
        self.improved_tab.setLayout(self.improved_layout)
        
        self.improved_code_edit = QTextEdit()
        self.improved_layout.addWidget(QLabel("改进后代码"))
        self.improved_layout.addWidget(self.improved_code_edit)
        
        # 添加所有标签页
        self.tab_widget.addTab(self.original_tab, "原始代码")
        self.tab_widget.addTab(self.requirements_tab, "需求说明")
        self.tab_widget.addTab(self.test_cases_tab, "测试用例")
        self.tab_widget.addTab(self.improved_tab, "改进后代码")
    
    def create_control_buttons(self):
        button_group = QGroupBox("控制")
        button_layout = QHBoxLayout()
        button_group.setLayout(button_layout)
        self.main_layout.addWidget(button_group)
        
        self.analyze_button = QPushButton("分析代码")
        self.analyze_button.clicked.connect(self.analyze_code)
        button_layout.addWidget(self.analyze_button)
        
        self.improve_button = QPushButton("实现改进")
        self.improve_button.clicked.connect(self.implement_improvements)
        button_layout.addWidget(self.improve_button)
        
        self.test_button = QPushButton("测试验证")
        self.test_button.clicked.connect(self.test_and_validate)
        button_layout.addWidget(self.test_button)
        
        self.iterate_button = QPushButton("完整迭代")
        self.iterate_button.clicked.connect(self.full_iteration)
        button_layout.addWidget(self.iterate_button)
        
        self.clear_button = QPushButton("清除日志")
        self.clear_button.clicked.connect(self.clear_logs)
        button_layout.addWidget(self.clear_button)
    
    def create_log_output(self):
        log_group = QGroupBox("日志输出")
        log_layout = QVBoxLayout()
        log_group.setLayout(log_layout)
        self.main_layout.addWidget(log_group)
        
        self.log_edit = QTextEdit()
        self.log_edit.setReadOnly(True)
        log_layout.addWidget(self.log_edit)
    
    # 以下是原有功能的GUI封装
    def load_api_key(self):
        key_path = Path.home() / ".deepseek_api_key"
        try:
            with open(key_path, 'r') as f:
                return f.read().strip()
        except FileNotFoundError:
            return None
    
    def save_api_key(self):
        key_path = Path.home() / ".deepseek_api_key"
        with open(key_path, 'w') as f:
            f.write(self.api_key_input.text())
        self.log_message("API Key已保存")
    
    def call_deepseek_api(self, messages, temperature=0.7, max_tokens=2000):
        headers = {
            "Authorization": f"Bearer {self.api_key_input.text()}",
            "Content-Type": "application/json"
        }
        
        payload = {
            "model": self.model_combo.currentText(),
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "stream": False
        }
        
        try:
            response = requests.post(
                "https://api.deepseek.com/v1/chat/completions",
                headers=headers,
                json=payload
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            self.log_message(f"API调用失败: {e}")
            return None
    
    def analyze_code(self):
        code = self.original_code_edit.toPlainText().strip()
        requirements = self.requirements_edit.toPlainText().strip()
        
        if not code:
            QMessageBox.critical(self, "错误", "请输入要分析的代码")
            return
        
        prompt = f"""
        你是一个高级代码分析AI。请分析以下代码并提出具体的改进建议。
        代码需要满足以下要求: {requirements}
        
        待分析代码:
        {code}
        
        请按照以下格式返回分析结果:
        1. 代码质量评估(1-10分)
        2. 主要问题列表
        3. 具体改进建议
        4. 优化后的完整代码(如果需要重大修改)
        """
        
        messages = [
            {"role": "system", "content": "你是一个专业的软件开发助手，擅长代码优化和重构。"},
            {"role": "user", "content": prompt}
        ]
        
        self.log_message("正在分析代码...")
        response = self.call_deepseek_api(messages)
        if response:
            result = response['choices'][0]['message']['content']
            self.log_message("分析结果:\n" + result)
            self.improved_code_edit.setPlainText(result)
        else:
            self.log_message("分析代码失败")
    
    def implement_improvements(self):
        original_code = self.original_code_edit.toPlainText().strip()
        suggestions = self.improved_code_edit.toPlainText().strip()
        
        if not original_code or not suggestions:
            QMessageBox.critical(self, "错误", "需要原始代码和改进建议")
            return
        
        prompt = f"""
        原始代码:
        {original_code}
        
        改进建议:
        {suggestions}
        
        请根据以上建议生成优化后的完整代码。保持原有功能不变，只进行质量改进。
        只返回最终的优化代码，不要包含任何解释。
        """
        
        messages = [
            {"role": "system", "content": "你是一个专业的软件开发助手，能够精确实现代码改进。"},
            {"role": "user", "content": prompt}
        ]
        
        self.log_message("正在实现改进...")
        response = self.call_deepseek_api(messages)
        if response:
            improved_code = response['choices'][0]['message']['content']
            self.log_message("改进后的代码已生成")
            self.improved_code_edit.setPlainText(improved_code)
        else:
            self.log_message("生成改进代码失败")
    
    def test_and_validate(self):
        code = self.improved_code_edit.toPlainText().strip()
        test_cases = self.test_cases_edit.toPlainText().strip()
        
        if not code:
            QMessageBox.critical(self, "错误", "请输入要测试的代码")
            return
        
        prompt = f"""
        请为以下代码设计全面的测试用例并执行测试:
        
        代码:
        {code}
        
        现有测试用例:
        {test_cases or '无'}
        
        请返回:
        1. 新增的测试用例(如有必要)
        2. 测试执行结果
        3. 发现的任何问题
        4. 修复建议(如有问题)
        """
        
        messages = [
            {"role": "system", "content": "你是一个专业的软件测试工程师，擅长设计全面的测试用例。"},
            {"role": "user", "content": prompt}
        ]
        
        self.log_message("正在运行测试...")
        response = self.call_deepseek_api(messages)
        if response:
            test_results = response['choices'][0]['message']['content']
            self.log_message("测试结果:\n" + test_results)
        else:
            self.log_message("测试验证失败")
    
    def full_iteration(self):
        original_code = self.original_code_edit.toPlainText().strip()
        requirements = self.requirements_edit.toPlainText().strip()
        test_cases = self.test_cases_edit.toPlainText().strip()
        
        if not original_code:
            QMessageBox.critical(self, "错误", "请输入原始代码")
            return
        
        self.log_message("开始完整迭代周期...")
        current_code = original_code
        for cycle in range(3):  # 最多3次迭代
            self.log_message(f"\n=== 迭代周期 {cycle + 1}/3 ===")
            
            # 分析代码
            self.log_message("分析代码中...")
            analysis = self.analyze_code()
            
            # 实现改进
            self.log_message("实现改进中...")
            improved_code = self.implement_improvements()
            if improved_code:
                current_code = improved_code
            else:
                break
                
            # 测试验证
            self.log_message("运行测试中...")
            test_results = self.test_and_validate()
            
            # 检查是否达到满意状态
            if "没有发现问题" in test_results or "所有测试通过" in test_results:
                self.log_message("\n代码已达到满意状态!")
                break
    
    def log_message(self, message):
        self.log_edit.append(message)
        self.status_bar.showMessage(message)
    
    def clear_logs(self):
        self.log_edit.clear()

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = CodeImproverGUI()
    window.show()
    sys.exit(app.exec_())